#pragma once

#include "FVOperator.h"

/**
 * @file MultigridSolver.h
 * @author xingyifan
 * @date 2022-08-30 21:56
 *
 * @description: 多重网格法
 */

template <int Dim>
class MultigridSolver
{
public:
    MultigridSolver(std::string type = "Laplacian", Real coeff = -1, GhostFiller<Dim> ghostfiller = GhostFiller<Dim>())
        : type_(type), coeff_(coeff), ghostfiller_(ghostfiller) {}

    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-08-29 11:59
     *
     * @description: 求解函数
     */

    void solve(Tensor<Dim> &f, Tensor<Dim> &v, int nu1 = 4, int nu2 = 4, int times = 100)
    {
        Tensor<Dim> tv;

        // 循环求解，直到达到最大循环次数 times 或者相邻两次误差范数足够小
        do
        {
            tv = v;
            VC(f, v, nu1, nu2);
            times--;
        } while (TensorOp::grid(tv - v) / TensorOp::grid(v) > 1e-6 && times > 0);
    }

private:
    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-07-26 14:09
     *
     * @description: V 循环
     */

    void VC(Tensor<Dim> &f, Tensor<Dim> &v, int nu1, int nu2)
    {
        // 网格分割
        int n = f.width() - 1;

        // 进行 nu1 次松弛
        GS(f, v, nu1);

        // laplacian
        Laplacian<Dim> laplacian;

        // 如果不是最粗的网格，由于有限体积的宽度为 5，不能太粗
        if (n > 4)
        {
            // 获得残量
            Tensor<Dim> dv;

            if (type_ == "I-Laplacian")
            {
                // 迭代矩阵是 I-k nu gamma L
                dv = TensorOp::add(f - v, TensorOp::multiple(laplacian(v, ghostfiller_), coeff_));
            }
            else
            {
                // 迭代矩阵是 L，注意 f 还要用，不能直接用 TensorOp::sub
                dv = f - laplacian(v, ghostfiller_);
            }

            // 约束算子
            Tensor<Dim> f2 = op_.FullWeighting(dv);

            // 粗网格上的初值
            Tensor<Dim> v2(n / 2 + 1, 0);

            // 在加粗网格上 V 循环
            VC(f2, v2, nu1, nu2);

            // 插值校正
            v += op_.Linear(v2);
        }

        // 进行 nu2 次松弛
        GS(f, v, nu2);
    }

    /**
     * @file MultigridSolver.h
     * @author xingyifan
     * @date 2022-08-31 22:50
     *
     * @description: GS 迭代法
     */

    void GS(Tensor<Dim> &f, Tensor<Dim> &v, int times)
    {
        // 加权系数
        Real omega = 2.0 / 3;

        // 网格规模
        int n = f.width();

        // 每段长度
        Real h = 1.0 / n;

        // 中心系数，它是 A[i,i] 元素的值
        Real center;

        if (type_ == "I-Laplacian")
        {
            // I-k nu gamma L
            center = 1 + coeff_ * 30 * Dim / (12 * h * h);
        }
        else
        {
            // Laplacian
            center = -30 * Dim / (12 * h * h);
        }

        // 加权循环
        while (times > 0)
        {
            // 遍历每一个点
            do
            {
                // 填充 ghost cell
                Vector ghost = ghostfiller_(v);

                // 记录总和
                Real sum = 0;

                // 加权迭代
                for (int d = 0; d < Dim; d++)
                {
                    sum -= -coeff_ *
                           (-ghost[5 * d] + 16 * ghost[5 * d + 1] + 16 * ghost[5 * d + 3] - ghost[5 * d + 4]) /
                           (12 * h * h);
                }

                v.value() = (1 - omega) * v.value() + omega * (sum + f.value()) / center;
            } while (v++, f++);
            times--;
        }
    }

public:
    std::string type_;             // 系数矩阵类型
    Real coeff_;                   // 迭代系数
    FVGridOperator<Dim> op_;       // 网格算子
    GhostFiller<Dim> ghostfiller_; // 携带 ghost cell 填充器
};